{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64b6c94f-3e79-46b7-b116-7027966777f8",
   "metadata": {},
   "source": [
    "# Mamba-2 Language Model demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fa052505-d91c-4e87-8daa-2b00ad8cc881",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a74ee1cb-b4b2-46a8-98a4-1dd845c1e5ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import torch\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "from mamba2 import Mamba2LMHeadModel\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device('mps')\n",
    "else:\n",
    "    device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a59ab109-2cbe-4f7a-b5ce-a58b8860f98c",
   "metadata": {},
   "source": [
    "Official pretrained models on [huggingface](https://huggingface.co/state-spaces):\n",
    "* `state-spaces/mamba2-130m`\n",
    "* `state-spaces/mamba2-370m`\n",
    "* `state-spaces/mamba2-780m`\n",
    "* `state-spaces/mamba2-1.3b`\n",
    "* `state-spaces/mamba2-2.7b`\n",
    "\n",
    "Choose a model depending on available system RAM (for CPU or system with unified memory) or VRAM.\n",
    "\n",
    "Note that these are base models without fine-tuning for downstream tasks such as chat or instruction following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6569ffd-993f-4d5b-9094-902801fe6c14",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "model = Mamba2LMHeadModel.from_pretrained(\"state-spaces/mamba2-1.3b\", device=device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\")\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bb837263-8a1f-40bf-a9b1-fce72225a674",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_config = dict(\n",
    "    max_new_length=200,\n",
    "    temperature=1.0,\n",
    "    top_k=30,\n",
    "    top_p=1.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "87006a5d-7992-4026-9b40-36cbc3ebf8dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(prompt: str, seed: int = 0, show_perf: bool = True):\n",
    "    \"\"\"Generate streaming completion\"\"\"\n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "    input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)[0]\n",
    "    print(prompt, end=\"\")\n",
    "\n",
    "    start = time.process_time()\n",
    "    n_generated = 0\n",
    "    for i, (token_id, _hidden_state) in enumerate(model.generate(input_ids, **generation_config)):\n",
    "        token = tokenizer.decode([token_id])\n",
    "        if i == 0:\n",
    "            now = time.process_time()\n",
    "            prompt_eval_elapsed, start = now - start, now\n",
    "        else:\n",
    "            n_generated += 1\n",
    "        print(token, end=\"\", flush=True)\n",
    "    if show_perf:\n",
    "        elapsed = time.process_time() - start\n",
    "        print('\\n\\n---')\n",
    "        print(f'Prompt eval | tokens: {input_ids.shape[0]} | elapsed: {prompt_eval_elapsed:.2f}s | tok/s: {input_ids.shape[0] / prompt_eval_elapsed:.2f}')\n",
    "        print(f'Generation | tokens: {n_generated} | elapsed: {elapsed:.2f}s | tok/s: {n_generated / elapsed:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4b926b16-2883-4eef-9459-3718498409e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mamba is a new state space model architecture that enables the modeling of discrete events in humanoid robots with simple and intuitive syntax.\n",
      "\n",
      "The Mamba state model is based on the state space model architecture of the state machine.\n",
      "Mamba enables fast and intuitive specification of the state transitions, without requiring any experience with formal modeling.\n",
      "The states are described on a per-event basis and they are not tied to an explicit representation of the robot world or\n",
      "the physics of the physical robot.\n",
      "\n",
      "Mamba is a free and open-source state space model software.\n",
      "\n",
      "For information on Mamba, visit mamba-robots.org\n",
      "\n",
      "What is a state machine?\n",
      "\n",
      "State machine modeling was pioneered by J.R. Walker and his colleagues in the early 1960s at MIT, who showed that\n",
      "continuous-time systems can be well represented by a simple discrete state machine. They also used this idea to build\n",
      "the first model of the humanoid robotic system known as Quoogle. Over the\n",
      "\n",
      "---\n",
      "Prompt eval | tokens: 9 | elapsed: 1.11s | tok/s: 8.08\n",
      "Generation | tokens: 199 | elapsed: 12.94s | tok/s: 15.38\n"
     ]
    }
   ],
   "source": [
    "generate(\"Mamba is a new state space model architecture\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "608ccece-9a11-47bc-bafd-7b47fc6383c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The meaning of life is death. But there is always a possibility that people may believe the opposite, as many have in various parts of the world, such as in India. The idea of God being the one who decides everything and life meaning is not decided by our thoughts, but by events. Life is not a fairytale and even if death is the only real possibility that people do not think of.\n",
      "\n",
      "India is the birthplace of Hinduism and the country has a history of several ancient civilizations. But what has remained unknown is the fact that Hinduism was not a religious system to worship in the past. It was more of a system of beliefs to live a better life. The most important point that can be ascertained is that life is all about choice and free will. The one who chooses the path, chooses the future that will be his.\n",
      "\n",
      "The Hindu way of life has been influenced by ancient Hindu traditions and beliefs. While the major tenets remain the same, the practices and rituals have\n",
      "\n",
      "---\n",
      "Prompt eval | tokens: 5 | elapsed: 0.33s | tok/s: 14.95\n",
      "Generation | tokens: 199 | elapsed: 12.53s | tok/s: 15.88\n"
     ]
    }
   ],
   "source": [
    "generate(\"The meaning of life is\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f0bc8a2b-b2bc-4d30-bf4c-213baec7441a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA is Nvidia's biggest moat on graphics hardware, and it's one that the gaming and PC markets have both been fighting to maintain for the last decade. However, Nvidia's Pascal architecture is on the horizon. And that could be a big opportunity for AMD.\n",
      "\n",
      "\n",
      "When Nvidia first released its Turing architecture at GTC back in February it was only on the cards; we only got an early taste of it. And so it has taken AMD quite some time to start taking a look at all that Nvidia-Turing design. However, AMD still has plenty of time to get to Nvidia before it's too late, and it needs to keep its eye on the Pascal architecture.\n",
      "\n",
      "AMD is also working on a new GPU called Vega which is due to go live this year; the first Vega GPU is rumored to be a reimagined Polaris architecture that features an enhanced memory hierarchy which could significantly speed up the graphics pipeline. If Vega is anything like Polaris,\n",
      "\n",
      "---\n",
      "Prompt eval | tokens: 9 | elapsed: 0.57s | tok/s: 15.67\n",
      "Generation | tokens: 199 | elapsed: 12.67s | tok/s: 15.71\n"
     ]
    }
   ],
   "source": [
    "generate(\"CUDA is Nvidia's biggest moat\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
